package com.maker.mcp.client.utils;

import com.maker.mcp.client.enums.SSEMsgType;
import lombok.extern.slf4j.Slf4j;
import org.springframework.util.CollectionUtils;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Consumer;

@Slf4j
public class SSEserver {

    public static final Map<String, SseEmitter> sseClients = new ConcurrentHashMap<>();
    /**
     * 创建sse连接
     * @param userId
     * @return
     */
    public static SseEmitter connect(String userId) {

        //设置超时时间，0L表示无超时,默认30s,超时未完成，则会抛出异常
        SseEmitter emitter = new SseEmitter(0L);

        //注册回调方法
        emitter.onTimeout(timeoutCallback(userId));
        emitter.onCompletion(completionCallback(userId));
        emitter.onError(errorCallback(userId));


        sseClients.put(userId, emitter);

        log.info("用户{}SSE连接被成功创建",userId);
        return emitter;
    }

    public static Runnable timeoutCallback(String userId) {
        return () -> {
            log.info("SSE超时");
            //删除用户连接
            remove(userId);
        };
    }

    public static Consumer<Throwable> errorCallback(String userId) {
        return Throwable -> {
            log.info("SSE异常");
            //删除用户连接
            remove(userId);
        };
    }

    public static Runnable completionCallback(String userId) {
        return () -> {
            log.info("SSE完成");
            remove(userId);
        };
    }

    public static void sendMsg(String userId, String message, SSEMsgType  msgType){
        if(CollectionUtils.isEmpty(sseClients)){
            return;
        }
        if(sseClients.containsKey(userId)){
            SseEmitter sseEmitter = sseClients.get(userId);
            sendEmitterMessage(sseEmitter, userId, message, msgType);
        }
    }


    public static void sendMsgToAllUsers(String message){
        if(CollectionUtils.isEmpty(sseClients)){
            return;
        }
        sseClients.forEach((userId, sseEmitter) -> {
            sendEmitterMessage(sseEmitter, userId, message, SSEMsgType.MESSAGE);
        });
    }
    private static void sendEmitterMessage(SseEmitter sseEmitter, String userId, String message, SSEMsgType  msgType) {

        try {
            SseEmitter.SseEventBuilder event = SseEmitter.event()
                    .id(userId)
                    .data(message)
                    .name(msgType.type);
            sseEmitter.send(event);
        } catch (IOException e) {
            log.info("用户{}SSE发送异常:{}",userId,e.getMessage());
            remove(userId);
        }
    }

    public static void remove(String userId){
        sseClients.remove(userId);
        log.info("用户{}SSE连接被移除，断开连接",userId);
    }
}
